import os, sys
import us
import datetime
from dateutil.relativedelta import *
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import statsmodels.api as sm
from statsmodels.tsa.api import VAR
from numpy.linalg import LinAlgError
from itertools import product
import pandas_datareader as pdr

PATH = r'C:\Users\Jeff Levy\Desktop\replication'

def run_var(df, unc, outcome, signif=0.05, maxlags=6, mindate=None, maxdate=None, make_indv_states=False):
    """Main function."""
    data, endog = setup_data(df, unc, outcome, mindate, maxdate)

    irf_lines, irf_summary = [], ['State,Description,Overshoot,Peak,Peak t,Min,Min t,Return t,Initial dip\n']
    for state in data.index.get_level_values(0).unique():
        st = data.loc[state]
        st = st.asfreq('MS')

        try:
            model = VAR(endog=st[endog])
            results = model.fit(maxlags=maxlags, ic=None)

            irf_summary, irf_lines = generate_irfs(results, endog, signif, state, irf_lines, irf_summary, make_indv_states)
        except LinAlgError:
            print(f'{state} did not converge.')

    if make_indv_states: write_summary(irf_summary)
    generate_irf_cloud(irf_lines, outcome, unc)
    print(f'Completed.  Endog: {endog}')
    return data

def setup_data(df, unc, outcome, mindate, maxdate):
    endog = [unc, outcome]
    df = df.sort_values(['state', 'date'])

    if mindate:
        df = df[df.index.get_level_values(1) >= mindate]
    if maxdate:
        df = df[df.index.get_level_values(1) <= maxdate]

    df = df.dropna(subset=endog, how='any') #drops the lead values from the hamilton filter that become nan
    df = df[endog].copy()

    return df, endog

def generate_irfs(results, endog, signif, state, irf_lines, irf_summary, make_indv_states):
    irf = results.irf(60)
    fig = irf.plot(impulse=endog[0],
                   response=endog[1],
                   signif=signif,
                   orth=True,
                   figsize=(12,5))

    ax = fig.get_axes()[0]
    lines = ax.get_lines()

    irf_lines.append((state, lines[0])) #collecton for plot of all IRFs on one

    t, y = lines[0].get_data()
    peak_unemp = y.max()
    peak_date = t[np.where(y == y.max())[0][0]]
    dip_unemp = y.min()
    dip_date = t[np.where(y == y.min())[0][0]]

    irf_summary.append(f'{state},,,{peak_unemp},{peak_date},{dip_unemp},{dip_date},,\n')
    fig.suptitle('')
    ax.set_title(state)

    if make_indv_states:
        fname = state + '.png'
        opath = os.path.join(PATH, 'irfs', fname)
        print('Saving: ', opath)
        fig.savefig(opath)
    plt.close()
    return irf_summary, irf_lines

def generate_irf_cloud(irf_lines, outcome, unc):
    if 'hf' in unc:
        ylab = 'Hamilton Filter'
    elif 'hp' in unc:
        ylab = 'HP Filter'
    else:
        ylab = 'Unfiltered'
    n = len(irf_lines)

    fig, ax = plt.subplots(figsize=(12,5))
    for state, lines in irf_lines:
        t, y = lines.get_data()
        if state == 'National':
            ax.plot(t, y, color='k', linestyle='-', label='National')
            n -= 1
        else:
            ax.plot(t, y, color='grey', linestyle='-', alpha=0.2, label='')
    ax.axhline(0, color='k', linestyle='-', linewidth=0.75)
    ax.plot(np.NaN, np.NaN, label=f'States (n={n})', color='grey', linestyle='-', alpha=0.2)

    ax.set_ylabel(ylab, fontsize=17.28*1.5)
    ax.xaxis.set_tick_params(labelsize=14.4*1.5)
    ax.yaxis.set_tick_params(labelsize=14.4*1.5)
    fig.tight_layout()
    fname = f'{unc},{outcome}.png'
    opath = os.path.join(PATH, 'irfs', fname)
    print('Saving: ', opath)
    fig.savefig(opath)
    plt.close()

if __name__ == '__main__':
    df = pd.read_csv(os.path.join(PATH, 'xtvar_ready.csv'), parse_dates=True, index_col=['state', 'date'])
    outcome_var = 'hf_unemp_rate'
    uncertainty_var = 'hf_ln_EPU_Composite'

    data = run_var(df, unc=uncertainty_var,
                       outcome=outcome_var,
                       mindate=datetime.datetime(2001, 7, 1),
                       maxdate=datetime.datetime(2019, 12, 1)) #min and max date are inclusive
